﻿
# -*- coding: utf-8 -*-
"""
Created on Tue May 12 12:56:10 2020
参考代码：https://www.tensorflow.org/tutorials/quickstart/advanced
"""

# ---------------------------
# （1）数据读取模块实验6-2
# ---------------------------
import os
from tensorflow.keras.datasets.mnist import load_data


class DataSource():
    def __init__(self):
        # 指定MNIST数据集本地存储路径
        data_path = 'D:\Data\实验9： 用CNN实现手写字识别\实验9 用CNN实现手写字识别\mnist.npz'

        # 加载数据集（返回训练集和测试集的元组）
        (x_train, y_train), (x_test, y_test) = load_data(path=data_path)

        # 数据预处理：
        # 1. 将图像数据从(28,28)扩展为(28,28,1)的四维张量（通道维度）
        x_train = x_train[..., tf.newaxis]
        x_test = x_test[..., tf.newaxis]

        # 2. 将像素值归一化到0-1范围（提升训练稳定性）
        x_train, x_test = x_train / 255.0, x_test / 255.0

        # 保存处理后的数据和标签
        self.train_images, self.train_labels = x_train, y_train
        self.test_images, self.test_labels = x_test, y_test

# 测试数据加载（可取消注释进行快速验证）
# data = DataSource()
# ---------------------------
# （2）模型构建模块
# ---------------------------
import tensorflow as tf
from tensorflow.keras import layers, models


class CNN():
    def __init__(self):
        # 初始化Sequential模型（按顺序堆叠层）
        model = models.Sequential()

        # 第1层：卷积层
        # 参数说明：
        # 32 - 输出通道数（即卷积核数量）
        # (3,3) - 卷积核尺寸
        # activation='relu' - 使用ReLU激活函数
        # input_shape=(28,28,1) - 输入数据形状（单通道灰度图）
        model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))

        # 第1个池化层（最大池化）
        # (2,2) - 池化窗口尺寸（缩小图像尺寸）
        model.add(layers.MaxPooling2D((2, 2)))

        # 第2层：卷积层（增加特征提取能力）
        model.add(layers.Conv2D(64, (3, 3), activation='relu'))

        # 第2个池化层
        model.add(layers.MaxPooling2D((2, 2)))

        # 第3层：卷积层（更深层次的特征提取）
        model.add(layers.Conv2D(64, (3, 3), activation='relu'))

        # 展平层（将三维特征图转换为一维向量）
        model.add(layers.Flatten())

        # 全连接层（64个神经元，ReLU激活）
        model.add(layers.Dense(64, activation='relu'))

        # 输出层（10个神经元对应0-9数字，Softmax激活）
        model.add(layers.Dense(10, activation='softmax'))

        # 输出模型结构概览
        model.summary()

        # 保存模型实例
        self.model = model


# 测试模型构建（可取消注释查看模型结构）
# network = CNN()

# ---------------------------
# （3）模型训练模块
# ---------------------------
class Train:
    def __init__(self):
        # 初始化模型和数据对象
        self.network = CNN()
        self.data = DataSource()

    def train(self):
        # 设置模型检查点（保存最佳权重）
        check_path = 'D:/test/ckpt/cp-{epoch:04d}.ckpt'
        save_model_cb = tf.keras.callbacks.ModelCheckpoint(
            check_path,
            save_weights_only=True,  # 仅保存权重
            verbose=1,  # 输出保存日志
            save_freq='epoch' , # 每个epoch保存一次
            period=5                # 注释掉后默认每个epoch保存
        )

        # 编译模型（配置训练参数）
        self.network.model.compile(
            optimizer='adam',  # 使用Adam优化器
            loss='sparse_categorical_crossentropy',  # 多分类交叉熵损失
            metrics=['accuracy']  # 监控准确率
        )

        # 开始训练
        self.network.model.fit(
            self.data.train_images,  # 训练数据
            self.data.train_labels,  # 训练标签
            epochs=10,  # 训练轮次
            callbacks=[save_model_cb]  # 使用回调函数
        )

        # 在测试集上评估模型
        test_loss, test_acc = self.network.model.evaluate(
            self.data.test_images,
            self.data.test_labels
        )

        # 输出测试结果
        print("准确率:{0:.2f}%，共测试了{1}张图片 ".format(test_acc * 100, len(self.data.test_labels)))


if __name__ == "__main__":
    # 启动训练过程
    mnist_train = Train()
    mnist_train.train()